package com.vincent.template.config;

import com.google.common.util.concurrent.RateLimiter;
import com.vincent.template.commons.exception.Vct520Exception;
import org.apache.commons.lang3.StringUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author Vincent Wang
 * @date 2022/9/7 18:37
 */
@Component
@Aspect
public class RateLimitAspect {

    private final Map<String, RateLimiter> rateLimiterMap = new ConcurrentHashMap<>();

    @Around(value = "@annotation(com.vincent.template.config.RateLimit)")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        Method method = ((MethodSignature) point.getSignature()).getMethod();
        RateLimit rateLimit = method.getAnnotation(RateLimit.class);
        String key = StringUtils.defaultIfBlank(rateLimit.name(),
                method.getDeclaringClass().getSimpleName() + "#" + method.getName());
        RateLimiter limiter = rateLimiterMap.computeIfAbsent(key, k -> RateLimiter.create(rateLimit.qps()));
        if (limiter.tryAcquire()) {
            return point.proceed();
        }
        throw new Vct520Exception(rateLimit.msg());
    }

}
